if '__main__' == __name__:
    import os
    import numpy as np
    from PyCmpltrtok.common import sep, rand_color
    from PyCmpltrtok.common_pymongo import PymongoDict
    from PyCmpltrtok.common_opencv import imzoom_rect
    import sys
    import cv2 as cv
    from coco_to_monogo import DB

    class CocoVisualEx(Exception):
        pass

    def _main():
        IMG_DIR = '/home/asuspei/large_data/_very_large/coco2014/coco_val2014/val2014'

        id2cap = PymongoDict(DB, 'id2cap')
        id2name, id2wh = PymongoDict(DB, 'id2name'), PymongoDict(DB, 'id2wh')
        cate_id2cate, cate_id2cate_super = PymongoDict(DB, 'cate_id2cate'), PymongoDict(DB, 'cate_id2cate_super')
        id2ann_ids = PymongoDict(DB, 'id2ann_ids')  # image id to annotations ids
        ann_id2cate_id, ann_id2bbox, ann_id2seg = PymongoDict(DB, 'ann_id2cate_id'), PymongoDict(DB, 'ann_id2bbox'), PymongoDict(DB, 'ann_id2seg')
        ids = PymongoDict(DB, 'lists')['ids']

        W, H = 1500, 700
        xpos = 16
        TGT_RECT = (W, H)
        WINDOW_NAME = 'check coco'

        xlen = len(ids)
        try:
            while True:
                cv.destroyAllWindows()

                xid = ids[xpos]
                xname = id2name[xid]
                xpath = os.path.join(IMG_DIR, xname)
                xcap = id2cap[xid]
                sep(xname)
                print(f'Pos: {xpos}')
                print(f'Path: {xpath}')
                print(f'Caption: {xcap}')
                w, h = id2wh[xid]
                print(f'w, h = {w}, {h}')

                img_arr = []
                ann_ids = id2ann_ids[xid]
                img_ = cv.imread(xpath, cv.IMREAD_COLOR)
                ann_id2color = {}

                def draw_them(ann_id_hl=None):
                    # bbox
                    img = img_.copy()
                    idx = -1
                    for ann_id in ann_ids:
                        idx += 1
                        cate_id = ann_id2cate_id[ann_id]
                        label = cate_id2cate[cate_id]
                        cate_super = cate_id2cate_super[cate_id]
                        xshow = f'{label} ({cate_super})'
                        xy12 = ann_id2bbox[ann_id]
                        xy12 = np.round(xy12).astype(int)
                        x1, y1, w, h = xy12
                        x2, y2 = x1 + w, y1 + h
                        if ann_id_hl is None or ann_id_hl == ann_id:
                            print(f'idx = {idx}', f'ann_id = {ann_id}', xshow, [x1, y1, x2, y2])
                        color = rand_color()
                        ann_id2color[ann_id] = color
                        cv.rectangle(img, (x1, y1), (x2, y2), color, 1)
                        cv.putText(img, xshow, (x1, y1 + 15), cv.FONT_HERSHEY_PLAIN, 1, color, 1)
                    img_arr.append(img)

                    # segmentation
                    img = np.zeros_like(img_)
                    idx = -1
                    for ann_id in ann_ids:
                        idx += 1
                        segs = ann_id2seg[ann_id]
                        if isinstance(segs, dict):
                            segs = [segs['counts']]
                            if len(segs[0]) % 2:
                                if ann_id_hl is None or ann_id_hl == ann_id:
                                    xprompt = f'**E** idx = {idx}, ann_id = {ann_id}, segs = {len(segs_int)}, xid = {xid}, xname = {xname}'
                                    print(xprompt)
                                    print(xprompt, file=sys.stderr)
                                continue
                        segs_int = []
                        for seg in segs:
                            segs_int.append(np.array(seg, dtype=np.int32).reshape(-1, 2))
                        if ann_id_hl is None or ann_id_hl == ann_id:
                            print(f'idx = {idx}', f'ann_id = {ann_id}', f'segs = {len(segs_int)}')
                        color = ann_id2color[ann_id]
                        if ann_id_hl is not None and ann_id == ann_id_hl:
                            thickness = -1
                        else:
                            thickness = 1
                        cv.drawContours(img, segs_int, -1, color, thickness)
                    img_arr.append(img)

                    img = np.concatenate(img_arr, axis=1)
                    img = imzoom_rect(img, TGT_RECT)
                    cv.line(img, (W // 2, 1), (W // 2, H), (0, 255, 0), 2)
                    cv.imshow(WINDOW_NAME, img)
                    cv.moveWindow(WINDOW_NAME, 3200, 0)

                print('> Please press p/n/q/+/- on the GUI windows.')
                img_ori = imzoom_rect(img_, TGT_RECT)
                cv.imshow(xname, img_ori)
                cv.moveWindow(xname, 1600, 180)

                draw_them()
                print('> Press +/- on the GUI window to highlight instances in cycling:')
                ann_id_pos = None
                ann_id_len = len(ann_ids)
                while(True):
                    k = cv.waitKey()
                    if ord('p') == k:
                        xpos -= 1
                        xpos = xpos if xpos >= 0 else 0
                        break
                    elif ord('n') == k:
                        xpos += 1
                        xpos = xpos if xpos < xlen else xlen - 1
                        break
                    elif ord('q') == k:
                        print('Bye!')
                        raise CocoVisualEx()
                    elif ord('+') == k or ord('-') == k:
                        if ord('+') == k:
                            if ann_id_pos is None:
                                ann_id_pos = 0
                            else:
                                ann_id_pos += 1
                                ann_id_pos = ann_id_pos if ann_id_pos < ann_id_len else 0
                        else:
                            if ann_id_pos is None:
                                ann_id_pos = 0
                            else:
                                ann_id_pos -= 1
                                ann_id_pos = ann_id_pos if ann_id_pos >= 0 else ann_id_len - 1
                        cv.destroyWindow(WINDOW_NAME)
                        img_arr = []
                        draw_them(ann_ids[ann_id_pos])
        except CocoVisualEx as ex:
            cv.destroyAllWindows()

    _main()
